#include "splitByShape.h"

namespace sp::Model::Split{

void registe_splitByShape(){
    splitByShape::registe();
}

splitByShape::splitByShape(std::vector<std::string> v): PythonFunction("/home/lx/SmartPipe/src/core/functions/Model/Split/", "", "splitByShape", "splitByShape", v){
    assert(v.size() == 0);
    name = "splitByShape";
    class_name = "Model-Split-splitByShape";
    type = SCATTER;
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({BATCH_TENSOR},TENSOR));
}

splitByShape::splitByShape(): PythonFunction(){
    name = "splitByShape";
    class_name = "Model-Split-splitByShape";
    type = SCATTER;
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/Split/";
    pModuleHomePath = "";
    pModuleName = "splitByShape";
    pClassName = "splitByShape";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({BATCH_TENSOR},TENSOR));
}

splitByShape::~splitByShape(){

}

void splitByShape::registe(){
    Function::Register("Model-Split-splitByShape", createObject);
    Function::Register("Model-Split-splitByShape", createObject2);
}

Function* splitByShape::createObject(std::vector<std::string> params){
    return new splitByShape(params);
}

Function* splitByShape::createObject2(){
    return new splitByShape();
}

void splitByShape::start(){
    PythonFunction::defaultStart();
    pInstance = PyObject_CallObject(pConstruct, nullptr);
}

bool splitByShape::waitForResource(){
    return true;
}

/** 
    vector<vector<Batch_Tensor*>> -> [Batch_Tensor,Batch_Tensor,...]
        splitByShape: [Batch_Tensor,Batch_Tensor,...] -> [[Tensor,Tensor,...],[Tensor,Tensor,...],...]
    [[Tensor,Tensor,...],[Tensor,Tensor,...],...] -> vector<Tensor*> 
    
    vector<vector<Head*>> -> [[Batch_Tensor,Batch_Tensor,Batch_Tensor],[Batch_Tensor,Batch_Tensor,Batch_Tensor],...]
        splitByShape: [[Batch_Tensor,Batch_Tensor,Batch_Tensor],[Batch_Tensor,Batch_Tensor,Batch_Tensor],...] -> [[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],...]
    [[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],...] -> vector<Head*> (Head* -> Tensor* -> Tensor* -> ...)
**/
void splitByShape::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    assert(data_input.size() != 0);
    assert(data_input[0].size() == 1);
    // 输入的tensor转成python格式
    PyObject* data = convertToPython(data_input);
    // 调用python函数
    pArgs = PyTuple_New(1);
    PyTuple_SetItem(pArgs, 0, data);
    pReturn = PyObject_CallMethod(pInstance, "process", "O", pArgs);
    // 结果转换成cpp格式并添加到data_output
    if(data_input[0][0]->type == HEAD)
        convertToCpp(pReturn, data_input, data_output, 2);
    else
        convertToCpp(pReturn, data_input, data_output, 0);
    // 释放内存
    for(int i = 0; i < data_input.size(); i++){
        assert(data_input[i][0]->type == BATCH_TENSOR || (data_input[i][0]->type == HEAD && data_input[i][0]->getLength() == 3 && data_input[i][0]->next->type == BATCH_TENSOR && data_input[i][0]->next->next->type == BATCH_TENSOR && data_input[i][0]->next->next->next->type == BATCH_TENSOR));
        if(data_input[i][0]->subRefCount() == 0){
            executor->freeAll(data_input[i][0]);
        }
    }
}

bool splitByShape::releaseResource(){
    return true;
}

void splitByShape::finish(){
    PythonFunction::defaultFinish();
}

void splitByShape::copy(Function* other){
    Function::defaultCopy(other);
}

};